%=========================================================================%
% This program solves and estimates the parameters of the retirement 
% choice and asset accumulation model. There are three steps:
% 1. Obtain data moments (targets for estimation step)
% 2. Solve the model given a guess for the underlying parameters
% 3. Iterate step 2 by generating simulated versions of the data
%    moments and minimizing distance to the data moments. 
% 4. Obtain simulated retirement and saving histories in-simulation
%    using estimated parameters.
%=========================================================================%   

%Log on...
diary (sprintf('master_estimate_log_%s.txt',datestr(now,'yyyymmdd')));

%Code files (functions) are here
cd '/.../structural-model-code';
addpath ('...data/HRS MiCDA/Tables');
addpath ('...data/FSRDC/Tables');
addpath ('...data');
addpath ('.../structural-model-code/lininterp');

%-------------------------------------------------------------------------%
% Obtain data moments
%-------------------------------------------------------------------------%
% Treatment effect estimates for 56-64 y.o's starting from period 0 to 12
temp=csvread('te_emp_5664.csv',6,0,[6,0,18,1]);
TE_D=temp(1:end,1);
V_TE=temp(:,2).^2; %variances 
clear temp;

%LEHD control group employment rates
temp=csvread('Controllfp_levels.csv',1,0,[1,0,17,1]);
%Extract periods -5 to 12 
E_D=temp(:,1); 
V_E=temp(:,2).^2; %variances 
clear temp;

%Stack the moments (scale TE's)
scaler = 10;
dtarget{1,1} = [E_D;scaler.*TE_D];

%Use I as wt matrix  
dtarget{2,1} = eye(length(dtarget{1,1}));

%-------------------------------------------------------------------------%
% 2. Estimate parameters using SMM and  I as wt matrix 
% There are five parameters describing preferences
% i. sigma (EIS)
% ii. rho (AR(1) persistence of work disutility process)
% iii. sigma_v (AR(1) idiosyncratic variance of disutility process)
% iv. phi is the age trend slope
% v. gamma is the labor disutility constant
%-------------------------------------------------------------------------%

%Set seed
rng(1,'twister');

%-------------------------------------------------------------------------%
%Estimate parameters using first step weight matrix 
%-------------------------------------------------------------------------%
%Initialize parameters;
%sorted results of sobol sequence runs 
SMM_criteria = ...
    csvread('.../model-output/SMMcriteria_sobolseq_eis515_10wt.csv',0,0);
x0_sobol = round(SMM_criteria(13,1:5),4); 
f = @(p)smm(p,dtarget,scaler);
options = optimset('MaxIter',100,'Display','iter');
[x1,~,~,~] = fminsearch(f,x0_sobol,options)  
csvwrite('/shared/Patki_ra/Aging LFP/output/x1_sobolstrt.csv',x1);

%obtain policy functions for estimated model
x1 = csvread('/shared/Patki_ra/Aging LFP/output/x1_sobolstrt.csv');
[optfunc,optfunc_frz]=polfunc_function(x1);

%-------------------------------------------------------------------------%
% 3. Compute VCV of the simulated moments 
%-------------------------------------------------------------------------%
p_star1 = [x1(1);1/(1+exp(x1(2)));exp(x1(3));x1(4);x1(5)];
%Define some input parameters:
rho = p_star1(2);
sigma_v = p_star1(3);
vcv_params=[rho;sigma_v];
%Omega is VCV
[Sm,Omega]=vcv_matrix(500,30,optfunc,optfunc_frz,vcv_params,scaler); 

%-------------------------------------------------------------------------%
% 4. Estimate standard errors using numerical derivatives
% approximate derviate as:
% dm/dtheta ~= (m(theta1+h) - m(theta1))/(h)
% where theta1 is the parameter estimate
% set step size as .005 
%-------------------------------------------------------------------------%

theta1 = [x1(1);1/(1+exp(x1(2)));exp(x1(3));x1(4);x1(5)];
h = 0.005;

%Matrix of perturbed parameters
thetadev = repmat(theta1',length(theta1'),1);
for i = 1:length(theta1)
    thetadev(i,i) = theta1(i)+h;
end

%Re-scale for model input
x1dev = thetadev;
x1dev(:,2) = log(1./thetadev(:,2) -1);
x1dev(:,3) = log(thetadev(:,3));


%Moments evaluated at the parameter estimate (simulated moments)
[E_S,~,~,~,~,TE_S]=smom_function(x1); 
SimM = [E_S(:,1);scaler.*TE_S]; 

%Compute Jacobian matrix
J = zeros(length(SimM),length(theta1));
for i = 1:length(theta1)
    
    %Simulated moments at perturbed parameter value
    [E_S_p,~,~,~,~,TE_S_p]=smom_function(x1dev(i,:));
    SimM_theta_plus_h = [E_S_p(:,1);scaler.*TE_S_p]; 
    %Derivatives
    J(:,i) = (SimM_theta_plus_h - SimM)./h;
    
end

%Standard errors using different wt matrices
%Omega is the VCV of the simulated moments
%Simple wt matrix 
%(I with treatment effect moments scaled up by 10x, so Var scaled by 10^2x)
W_simp = eye(length(SimM));
for i = 18:length(SimM)
    W_simp(i,i) = scaler^2;
end

%Variance matrices and standard errors
Nsims = 5000; %number of simulations
V_robust = (1+1/Nsims).*inv(J'*W_simp*J)*(J'*W_simp*Omega*W_simp*J)*inv(J'*W_simp*J);
sevec_robust = sqrt(diag(V_robust));
V_optimal = (1+1/Nsims).*inv(J'/Omega*J);
sevec_optimal = sqrt(diag(V_optimal));

%-------------------------------------------------------------------------%
% 5. Chi-2 model-fit statistics
%-------------------------------------------------------------------------%
%Data moments
DatM = dtarget{1,1};

%SMM objective function (J-stat)
m = SimM - DatM;
J = m'/Omega*m;
df = length(SimM)-length(theta1);
chi2test = [J df];

%-------------------------------------------------------------------------%
% 6. Output the results
%-------------------------------------------------------------------------%
paramse = [theta1 sevec_robust];
csvwrite('.../model-output/chi2test_Iwt.csv',chi2test);
csvwrite('.../model-output/paramest_Iwt.csv',paramse);
%Simulated moments
csvwrite('.../model-output/E_S.csv',E_S);
csvwrite('.../model-output/TE_S.csv',TE_S);